from gmssl import sm4
import base64 as b64

class SM4():
    def __init__(self):
        self.gmsm4 = sm4.CryptSM4()  # 实例化

    def encryptSM4_StrtoBase64(self, encrypt_key, value):
        if type(encrypt_key) == str:
            encrypt_key =  encrypt_key.encode(encoding='UTF-8')
        self.gmsm4.set_key(encrypt_key, sm4.SM4_ENCRYPT)
        if type(value) != str:
            value = str(value)
        encrypt_value = self.gmsm4.crypt_ecb(value.encode())
        return b64.b64encode(encrypt_value).decode('UTF-8')

    def decryptSM4_Base64toStr(self, decrypt_key, value):
        if type(decrypt_key) == str:
            decrypt_key = decrypt_key.encode(encoding='UTF-8')
        self.gmsm4.set_key(decrypt_key, sm4.SM4_DECRYPT)
        tmpBytes = b64.b64decode(value)
        decrypt_value = self.gmsm4.crypt_ecb(tmpBytes).decode(encoding='UTF-8')
        return decrypt_value



